<?php

namespace app\common\library;

use think\facade\Cache;

/**
 * 基于Redis使用令牌桶算法实现流量控制
 * 使用redis的队列作为令牌桶容器，入队（lPush）出队（rPop)作为令牌的加入与消耗操作。
 * Class Shaper
 * @package app\common\library
 */
class Shaper
{
    protected string $container;  // 令牌桶容器
    protected int $maxToken;      // 最大令牌数
    protected object $redis;      // Redis对象

    private int $selectDb = 15;   // 选择存储的数据库

    /**
     * 初始化
     * @param string $container 令牌桶容器
     * @param int $maxToken 最大令牌数
     */
    public function __construct(string $container, int $maxToken = 100)
    {
        $this->container = $container;
        $this->maxToken = $maxToken;
        $this->redis = $this->connectRedis();
    }

    /**
     * 加入令牌
     * @param int $num 加入的令牌数量
     * @return int 加入的数量
     */
    public function add(int $num = 0): int
    {
        // 剩余令牌数
        $curNum = intval($this->redis->lLen($this->container));
        // 最大令牌数
        $maxNum = intval($this->maxToken);
        // 计算最大可加入的令牌数量，不能超过最大令牌数
        $num = $maxNum >= $curNum + $num ? $num : $maxNum - $curNum;
        // 加入令牌
        if ($num > 0) {
            $token = array_fill(0, $num, 1);
            $this->redis->lPush($this->container, ...$token);
            return $num;
        }
        return 0;
    }

    /**
     * 获取令牌
     * @return bool
     */
    public function get(): bool
    {
        return $this->redis->rPop($this->container) ? true : false;
    }

    /**
     * 重设令牌桶，填满令牌
     */
    public function resetBucketToken()
    {
        $this->redis->del($this->container);
        $this->add($this->maxToken);
    }

    /**
     * 连接Redis，返回缓存对象
     * @return object
     */
    public function connectRedis(): object
    {
        $redis = Cache::store('redis')->handler();
        $redis->select($this->selectDb);

        return $redis;
    }
}